from Agent.Agent import Agent
import pandas as pd
from llm import llm

PATH = './rag_data/设备参数详情.xlsx'


def dict2str(dic):
    doc = ''
    for k, v in dic.items():
        if str(v) and str(v) != 'nan':
            doc += f'## {k}\n{v}\n'
    return doc


data1 = pd.read_excel(PATH)

data1 = data1.to_dict('records')
doc1 = []
for i in data1:
    doc1.append(dict2str(i))

data2 = pd.read_excel(PATH, sheet_name='字段释义')

data2 = data2.to_dict('records')

doc2 = ''
for i in data2:
    doc2 += dict2str(i)

prompt = '''
以下是设备参数运转文档：
---
{}
---
以下是字段释义：
---
{}
---
请基于文档回答问题，回答问题时尽可能和字段对应，确保和字段完全一致。
如果有缺失值，说明缺失后推理出合理答案。

# 关于设备参数详情的问题
a）上下限问题：有上限回答上限，有下限回答下限；
b）触发机制：优先考虑“报警/停机/安全保护/Alarm/Shutdown/Safety protection”

示例：
user:停泊发电机组转速为1800RPM会发生什么？
assistant:根据文档停泊发电机组转速上限为2000RPM超过75%(1500RPM)会触发`Alarm state`,答案:`Alarm`


'''


class RagAgent(Agent):
    def __init__(self, name,
                 description,
                 system_prompt,
                 agent_type,
                 llm_engine,
                 agents=None,
                 rag=None):
        super().__init__(name,
                         description,
                         system_prompt,
                         agent_type,
                         llm_engine,
                         agents,
                         rag)

    def build_prompt(self, question):
        self.messages = [{"role": "system", "content": self.system_prompt.format('\n'.join(doc1), doc2)},
                         {"role": "user", "content": question}]

    def run(self, question):
        self.build_prompt(question)
        return self.llm_engine(self.messages)


RagAgentConfig = {"name": 'xxx',
                  "description": 'xxx',
                  'system_prompt': prompt,
                  'agent_type': 'plan-action',
                  'llm_engine': llm,
                  'rag': None,
                  'agents': None}

rag_agent = RagAgent(**RagAgentConfig)

